// local
#include "TrackerKCF.hpp"
#include "FFTTools.hpp"
#include "RectTools.hpp"

#include <iostream>
#include <fstream>

using namespace std;

namespace handtrack {

//////////////
/** Public **/
//////////////
TrackerKCF::TrackerKCF(int patch_size, bool multiscale)
    : _patch_size(patch_size),
      _sigma(0.2f),
      _output_sigma_factor(0.125f),
      _lambda(0.0001f),
      _padding(3.f),
      _scale_step(1.f),
      _scale_weight(0.f) { 
    if (multiscale) { 
        this->_scale_step = 1.1f;
        this->_scale_weight = 0.9f;
    }

    this->createHanningMats();
    this->createGaussianPeak();
    this->_alphaf = cv::Mat::zeros(
        this->_patch_size,
        this->_patch_size,
        CV_32FC2
    );
    this->_tmpl = cv::Mat::zeros(
        this->_patch_size,
        this->_patch_size,
        CV_32F
    );
}

TrackerKCF::~TrackerKCF()
{
}

void TrackerKCF::update(
    const cv::Rect& src_roi,
    cv::Mat src_image,
    float train_interp_factor
)
{
    // printf("=================update=================\n");
    assert(src_roi.width > 0 && src_roi.height > 0);
    assert(train_interp_factor >= 0.f && train_interp_factor <= 1.f);

    this->adjustROI(src_roi, src_image);
    cv::Mat x = this->getFeatures(src_image, 1.f);
    cv::Mat k = this->gaussianCorrelation(x, x);

    cv::Mat kk = FFTTools::fftd(k);
    PRINT_CVMATComplex("kk", kk, 0, 0, 8, 8)

    kk = FFTTools::fftd(k)+this->_lambda;

    // printf("---------lambda:%f\n", this->_lambda);
    // cv::Mat alphaf = FFTTools::complexDivision(
    //     this->_prob,
    //     (FFTTools::fftd(k) + this->_lambda)
    // );
    cv::Mat alphaf = FFTTools::complexDivision(
        this->_prob,
        (FFTTools::fftd(k) + this->_lambda)
    );
    PRINT_CVMATComplex("alphaf", alphaf, 0, 0, 8, 8)
    this->_tmpl = (1.f - train_interp_factor) * this->_tmpl + train_interp_factor * x;
    this->_alphaf = (1.f - train_interp_factor) * this->_alphaf + train_interp_factor * alphaf;
    PRINT_CVMATF32("_tmpl", this->_tmpl, 0, 0, 8, 8)
    PRINT_CVMATComplex("_alphaf", this->_alphaf, 0, 0, 8, 8)
}

cv::Rect TrackerKCF::detect(
    const cv::Rect& dst_roi,
    cv::Mat dst_image,
    float& peak_value
)
{
    assert(dst_roi.width > 0 && dst_roi.height > 0);

    // printf("=================detect=================\n");
    this->adjustROI(dst_roi, dst_image);
    float cx = this->_roi.x + this->_roi.width / 2.f;
    float cy = this->_roi.y + this->_roi.height / 2.f;

    float res_peak_value = 0.f;
    cv::Point2f res = this->findPixelPeak(
        this->_tmpl,
        this->getFeatures(dst_image, 1.f),
        res_peak_value
    );
    peak_value = res_peak_value;

    if (this->_scale_step != 1.f)
    {
        // Test at a smaller scale
        res_peak_value = 0.f;
        cv::Point2f new_res = this->findPixelPeak(
            this->_tmpl,
            this->getFeatures(dst_image, 1.f / this->_scale_step),
            res_peak_value
        );

        if (this->_scale_weight * res_peak_value > peak_value) {
            res = new_res;
            peak_value = this->_scale_weight * res_peak_value;
            this->_scale /= this->_scale_step;
            this->_roi.width /= this->_scale_step;
            this->_roi.height /= this->_scale_step;
        }

        // Test at a bigger scale
        res_peak_value = 0.f;
        new_res = this->findPixelPeak(
            this->_tmpl,
            this->getFeatures(dst_image, this->_scale_step),
            res_peak_value
        );

        if (this->_scale_weight * res_peak_value > peak_value) {
            res = new_res;
            peak_value = this->_scale_weight * res_peak_value;
            this->_scale *= this->_scale_step;
            this->_roi.width *= this->_scale_step;
            this->_roi.height *= this->_scale_step;
        }
    }

    // printf("--------cx=%f,cy=%f,_roi.width=%f,_roi.height=%f,res.x=%f,res.y=%f\n", cx, cy, _roi.width, _roi.height, res.x, res.y);
    this->_roi.x = cx - this->_roi.width / 2.f + ((float) res.x * this->_scale);
    this->_roi.y = cy - this->_roi.height / 2.f + ((float) res.y * this->_scale);
    // printf("---------detect: [%f,%f,%f,%f], peak_value:%f\n", _roi.x, _roi.y, _roi.width, _roi.height, peak_value);
    return _roi;
}


/////////////////
/** Protected **/
/////////////////
cv::Point2f TrackerKCF::findPixelPeak(cv::Mat z, cv::Mat x, float& peak_value)
{
    // PRINT_CVMATF32("pixelpeak z", z)
    // PRINT_CVMATF32("pixelpeak x", x)
    cv::Mat k = this->gaussianCorrelation(x, z);
    cv::Mat res = FFTTools::real(
        FFTTools::fftd(
            FFTTools::complexMultiplication(
                this->_alphaf,
                FFTTools::fftd(k)
            ),
            true
        )
    );
    PRINT_CVMATF32("pixelpeak res", res, this->_patch_size/2, this->_patch_size/2, 8, 8)
    //minMaxLoc only accepts doubles for the peak, and integer points for the coordinates
    cv::Point2i pi;
    double pv;
    cv::minMaxLoc(res, NULL, &pv, NULL, &pi);
    peak_value = (float)pv;

    printf("-------[%f,],[%ld,%ld], \n",  pv,  pi.x, pi.y);
    //subpixel peak estimation, coordinates will be non-integer
    cv::Point2f p((float)pi.x, (float)pi.y);

    if (pi.x > 0 && pi.x < res.cols-1) {
        p.x += subPixelPeak(res.at<float>(pi.y, pi.x-1), peak_value, res.at<float>(pi.y, pi.x+1));
    }

    if (pi.y > 0 && pi.y < res.rows-1) {
        p.y += subPixelPeak(res.at<float>(pi.y-1, pi.x), peak_value, res.at<float>(pi.y+1, pi.x));
    }

    p.x -= (res.cols) / 2;
    p.y -= (res.rows) / 2;
    printf("------p=[%f,%f], %d,%d\n",p.x, p.y, res.cols, res.rows);
    return p;
}

void TrackerKCF::adjustROI(const cv::Rect& roi, cv::Mat image)
{
    float orig_cx = static_cast<float>(roi.x) + static_cast<float>(roi.width) / 2.f;
    float orig_cy = static_cast<float>(roi.y) + static_cast<float>(roi.height) / 2.f;
    if (roi.width >= roi.height)
    {
        this->_roi.width = static_cast<float>(roi.width);
        this->_roi.height = static_cast<float>(roi.width);
    }
    else
    {
        this->_roi.width = static_cast<float>(roi.height);
        this->_roi.height = static_cast<float>(roi.height);
    }
    this->_roi.x = orig_cx - this->_roi.width / 2.f;
    this->_roi.y = orig_cy - this->_roi.height / 2.f;
    if (this->_roi.x + this->_roi.width <= 0) this->_roi.x = -this->_roi.width + 1;
    if (this->_roi.y + this->_roi.height <= 0) this->_roi.y = -this->_roi.height + 1;
    if (this->_roi.x >= image.cols - 1) this->_roi.x = image.cols - 2;
    if (this->_roi.y >= image.rows - 1) this->_roi.y = image.rows - 2;

    float pad_cx = this->_roi.x + this->_roi.width / 2.f;
    float pad_cy = this->_roi.y + this->_roi.height / 2.f;
    float pad_length = min<float>(
        this->_roi.width * this->_padding,
        min<float>((float)image.rows, (float)image.cols)
    );
    this->_pad_roi.width = pad_length;
    this->_pad_roi.height = pad_length;
    this->_pad_roi.x = pad_cx - pad_length / 2.f;
    this->_pad_roi.y = pad_cy - pad_length / 2.f;

    this->_scale = pad_length / static_cast<float>(this->_patch_size);

    // printf("_roi:[%f,%f,%f,%f],_pad_roi:[%f,%f,%f,%f], _scale=%f\n", _roi.x, _roi.y, _roi.width, _roi.height, _pad_roi.x, _pad_roi.y, _pad_roi.width, _pad_roi.height, _scale);
}

cv::Mat TrackerKCF::gaussianCorrelation(cv::Mat x1, cv::Mat x2)
{
    // PRINT_CVMATF32("x1", x1);
    // PRINT_CVMATF32("x2", x2);
    cv::Mat c = cv::Mat::zeros(cv::Size(this->_patch_size, this->_patch_size), CV_32F);
    // cv::mulSpectrums(FFTTools::fftd(x1), FFTTools::fftd(x2), c, 0, true);
    cv::Mat x1_c = FFTTools::fftd(x1);
    cv::Mat x2_c = FFTTools::fftd(x2);
    PRINT_CVMATComplex("x1_c", x1_c, 0, 0, 8, 8);
    cv::mulSpectrums(x1_c, x2_c, c, 0, true);
    PRINT_CVMATComplex("spectrum", c, 0, 0, 8, 8);
    PRINT_CVMATComplex("spectrum", c, _patch_size/2, _patch_size/2, 8, 8);
    c = FFTTools::fftd(c, true);
    // PRINT_CVMATComplex("c inv", c);
    FFTTools::rearrange(c);
    c = FFTTools::real(c);
    PRINT_CVMATF32("c_real", c, _patch_size/2, _patch_size/2, 8, 8);
    float patch_step = static_cast<float>(this->_patch_size * this->_patch_size);
    cv::Mat d; 
    // cv::max(((cv::sum(x1.mul(x1))[0] + cv::sum(x2.mul(x2))[0]) - 2.f * c) / patch_step, 0, d);
    float x1_sum = cv::sum(x1.mul(x1))[0];
    float x2_sum = cv::sum(x2.mul(x2))[0];
    printf("------x1_sum=%f, x2_sum=%f, _sigma=%f\n", x1_sum, x2_sum, this->_sigma);
    cv::max(((x1_sum + x2_sum) - 2.f * c) / patch_step, 0, d);
    cv::Mat k;
    cv::exp((-d / (this->_sigma * this->_sigma)), k);
    // {
    //     float* buf = (float*)malloc(128*128*sizeof(float));
    //     char dbg_file_name[1024];
    //     sprintf(dbg_file_name, "/vendor/kcf/cpu_out/dbg_data.txt");
    //     std::ofstream dbg_out(dbg_file_name, std::ios::out|std::ios::trunc);
    //     for (int i = 0; i < 128; i++) {
    //         for (int j = 0; j < 128; j++) {
    //             dbg_out << k.at<float>(i, j) << ", "; 
    //         }
    //         dbg_out << endl;
    //     }
    //     free(buf);
    // }
    PRINT_CVMATF32("corr", k, _patch_size/2, _patch_size/2, 8, 8);
    return k;
}

cv::Mat TrackerKCF::getFeatures(const cv::Mat& image, float scale_adjust)
{
    int cx = static_cast<int>(this->_pad_roi.x + this->_pad_roi.width / 2.f);
    int cy = static_cast<int>(this->_pad_roi.y + this->_pad_roi.height / 2.f);
    cv::Rect extracted_roi;
    extracted_roi.width = static_cast<int>(this->_pad_roi.width * scale_adjust);
    extracted_roi.height = static_cast<int>(this->_pad_roi.height * scale_adjust);
    extracted_roi.x = cx - extracted_roi.width / 2;
    extracted_roi.y = cy - extracted_roi.height / 2;

    // printf("-------extracted_roi:[%d,%d,%d,%d]\n", extracted_roi.x, extracted_roi.y, extracted_roi.width, extracted_roi.height); 
    cv::Mat z = RectTools::subwindow(image, extracted_roi, cv::BORDER_REPLICATE);
    if (z.cols != this->_patch_size || z.rows != this->_patch_size) {
        // printf("-------------[%d,%d]---resize-----\n", z.rows, z.cols);
        cv::resize(z, z, cv::Size(this->_patch_size, this->_patch_size));
    } 
    // PRINT_CVMATU8("z", z);
    cv::Mat FeaturesMap = RectTools::getGrayImage(z);
    FeaturesMap -= 0.5f;
    FeaturesMap = this->_hann.mul(FeaturesMap);
    PRINT_CVMATF32("FeaturesMap", FeaturesMap, _patch_size/2, _patch_size/2, 8, 8)
    return FeaturesMap;
}

float TrackerKCF::subPixelPeak(float left, float center, float right)
{
    float divisor = 2.f * center - right - left;

    if (divisor == 0.f)
        return 0.f;
    
    return 0.5f * (right - left) / divisor;
}

void TrackerKCF::createHanningMats()
{
    cv::Mat hann1t = cv::Mat(cv::Size(this->_patch_size, 1), CV_32F, cv::Scalar(0));
    cv::Mat hann2t = cv::Mat(cv::Size(1, this->_patch_size), CV_32F, cv::Scalar(0)); 
    for (int i = 0; i < hann1t.cols; i++) {
        hann1t.at<float > (0, i) = 0.5 * (1 - std::cos(2 * 3.14159265358979323846 * i / (hann1t.cols - 1)));
    }
    for (int i = 0; i < hann2t.rows; i++) {
        hann2t.at<float > (i, 0) = 0.5 * (1 - std::cos(2 * 3.14159265358979323846 * i / (hann2t.rows - 1)));
    }

    this->_hann = hann2t * hann1t;
    PRINT_CVMATF32("hann", this->_hann, _patch_size/2, _patch_size/2, 8, 8)
}

void TrackerKCF::createGaussianPeak()
{
    cv::Mat_<float> res(this->_patch_size, this->_patch_size);

    int syh = this->_patch_size / 2;
    int sxh = this->_patch_size / 2;

    float output_sigma = std::sqrt((float) this->_patch_size * this->_patch_size) / this->_padding * this->_output_sigma_factor;
    float mult = -0.5 / (output_sigma * output_sigma);

    for (int i = 0; i < this->_patch_size; i++)
    {
        for (int j = 0; j < this->_patch_size; j++)
        {
            int ih = i - syh;
            int jh = j - sxh;
            res(i, j) = std::exp(mult * (float) (ih * ih + jh * jh));
        }
    }
    // PRINT_CVMATF32("res", res, _patch_size/2, _patch_size/2, 16, 16);
    this->_prob = FFTTools::fftd(res);
    PRINT_CVMATComplex("_prob", this->_prob, 0, 0, 8, 8);
    PRINT_CVMATComplex("_prob", this->_prob, _patch_size/2, _patch_size/2, 8, 8);
}

} // end namespace handtrack